package cn.ycc1.lianzi.controller;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Arrays;

/**
 * @author ycc
 * @date 2025/2/13
 */
@RestController
@RequestMapping("/api")
public class StreamController {
    // 流式调用接口
    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamGeneration(@RequestParam("prompt") String prompt)
            throws NoApiKeyException, ApiException, InputRequiredException {

        SseEmitter emitter = new SseEmitter(60_000L); // 60秒超时
        Generation gen = new Generation();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(prompt)
                .build();

        // 调用DashScope流式API
        GenerationParam param = buildGenerationParam(userMsg);
        Flowable<GenerationResult> resultFlowable = gen.streamCall(param);

        // 异步处理流
        resultFlowable.subscribe(
                message -> sendDataChunk(emitter, message),
                error -> handleError(emitter, error),
                () -> completeEmitter(emitter)
        );

        // 处理客户端断开连接
        emitter.onCompletion(() -> cleanupResources());
        emitter.onTimeout(() -> cleanupResources());

        return emitter;
    }

    // 流式调用接口
    @PostMapping(value = "/stream2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamGeneration2(@RequestBody UserRequest request)
            throws NoApiKeyException, ApiException, InputRequiredException {

        SseEmitter emitter = new SseEmitter(60_000L); // 60秒超时
        Generation gen = new Generation();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(request.getPrompt())
                .build();

        // 调用DashScope流式API
        GenerationParam param = buildGenerationParam(userMsg);
        Flowable<GenerationResult> resultFlowable = gen.streamCall(param);

        // 异步处理流
        resultFlowable.subscribe(
                message -> sendDataChunk(emitter, message),
                error -> handleError(emitter, error),
                () -> completeEmitter(emitter)
        );

        // 处理客户端断开连接
        emitter.onCompletion(() -> cleanupResources());
        emitter.onTimeout(() -> cleanupResources());

        return emitter;
    }

    // 发送单个数据块
    private void sendDataChunk(SseEmitter emitter, GenerationResult message) {
        try {
            String content = message.getOutput().getChoices().get(0).getMessage().getContent();
            emitter.send(SseEmitter.event().data(content)); // SSE格式
        } catch (Exception e) {
            handleError(emitter, e);
        }
    }

    // 错误处理
    private void handleError(SseEmitter emitter, Throwable error) {
        try {
            emitter.completeWithError(error);
        } catch (Exception e) {
            // 记录错误日志
        } finally {
            cleanupResources();
        }
    }

    // 完成发送
    private void completeEmitter(SseEmitter emitter) {
        try {
            emitter.complete();
        } catch (Exception e) {
            // 记录完成异常
        }
    }

    // 清理资源
    private void cleanupResources() {
        // 关闭数据库连接等操作（按需实现）
    }

    // 构建Generation参数
    private GenerationParam buildGenerationParam(Message userMsg) {
        return GenerationParam.builder()
                .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                .model("deepseek-r1")
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }

    // 请求参数类
    public static class UserRequest {
        private String prompt;

        // Getter/Setter
        public String getPrompt() { return prompt; }
        public void setPrompt(String prompt) { this.prompt = prompt; }
    }
}
